module weno_mod
  
  use constants_mod
  use namelist_mod

  implicit none
  
  private
  public compute_interface_flux_weno5
  public compute_flux_jacobian_weno5

  real(r8), parameter :: p1(3) = [ 1.d0 / 3.d0, -7.d0 / 6.d0, 11.d0 / 6.d0]
  real(r8), parameter :: p2(3) = [-1.d0 / 6.d0,  5.d0 / 6.d0,  1.d0 / 3.d0]
  real(r8), parameter :: p3(3) = [ 1.d0 / 3.d0,  5.d0 / 6.d0, -1.d0 / 6.d0]
  real(r8), parameter :: s1 = 13.d0 / 12.d0, s2 = 1.d0 / 4.d0
  real(r8), parameter :: s11(3) = [1.d0, -2.d0, 1.d0], s12(3) = [1.d0, -4.d0, 3.d0]
  real(r8), parameter :: s21(3) = [1.d0, -2.d0, 1.d0], s22(3) = [1.d0,  0.d0,-1.d0]
  real(r8), parameter :: s31(3) = [1.d0, -2.d0, 1.d0], s32(3) = [3.d0, -4.d0, 1.d0]
  real(r8), parameter :: gamma(3) = [1.d0 / 10.d0, 6.d0 / 10.d0, 3.d0 / 10.d0] ! linear weights

  contains

  subroutine compute_flux_jacobian_weno5(var, dt, jacobian)

    real(r8), intent(in) :: var(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    real(r8), intent(out) :: jacobian(1:nx, 1:nx)
    real(r8) :: dF_plus_du(-2:2), dF_minus_du(-2:2)
    integer :: i, j, k
    real(r8) alpha
    
    alpha = abs(u)
    do i = 1, nx
      ! Flux at i + 1/2
      call compute_dF_du(var, i  , alpha, dt, dF_plus_du)
      ! Flux at i - 1/2
      call compute_dF_du(var, i-1, alpha, dt, dF_minus_du)
      do k = -2, 2
        j = i + k
        if (j < 1 ) j = j + nx ! periodic condition
        if (j > nx) j = j - nx ! periodic condition
        jacobian(i,j) = -(dF_plus_du(k) - dF_minus_du(k)) / dx
      end do
    end do

  end subroutine compute_flux_jacobian_weno5

  subroutine compute_dF_du(var, i, alpha, dt, dF_du)

    real(r8), intent(in) :: var(1-nhalo:nx+nhalo), alpha, dt
    integer, intent(in) :: i
    real(r8), intent(out) :: dF_du(-2:2)
    real(r8) :: dwgt_plus_du(1:3,-2:2), dwgt_minus_du(1:3,-2:2)
    real(r8) :: du_plus_du(-2:2), du_minus_du(-2:2)
    real(r8) :: wgts(1:3), poly(3)
    integer :: k, m, j

    ! Left-biased WENO (u_{i+1/2}^-)
    call compute_weight_derivative(var, i, wgts, dwgt_minus_du, biased='left_stencil')
    poly(1) = dot_product(p1, var(i-2:i  ))
    poly(2) = dot_product(p2, var(i-1:i+1))
    poly(3) = dot_product(p3, var(i  :i+2))
    du_minus_du = 0.0_r8
    do k = 1, 3
      do m = -2, 2
        j = i + m
        if (j < 1) j = j + nx
        if (j > nx) j = j - nx
        du_minus_du(m) = du_minus_du(m) + wgts(k) * (&
            merge(p1(1), 0.0_r8, k==1 .and. m==-2) + & ! ∂p1/∂u_{i-2}
            merge(p1(2), 0.0_r8, k==1 .and. m==-1) + & ! ∂p1/∂u_{i-1}
            merge(p1(3), 0.0_r8, k==1 .and. m== 0) + & ! ∂p1/∂u_{i}
            merge(p2(1), 0.0_r8, k==2 .and. m==-1) + & ! ∂p2/∂u_{i-1}
            merge(p2(2), 0.0_r8, k==2 .and. m==0 ) + & ! ∂p2/∂u_{i}
            merge(p2(3), 0.0_r8, k==2 .and. m==1 ) + & ! ∂p2/∂u_{i+1}
            merge(p3(1), 0.0_r8, k==3 .and. m==0 ) + & ! ∂p3/∂u_{i}
            merge(p3(2), 0.0_r8, k==3 .and. m==1 ) + & ! ∂p3/∂u_{i+1}
            merge(p3(3), 0.0_r8, k==3 .and. m==2 )   & ! ∂p3/∂u_{i+2}
            ) + dwgt_minus_du(k,m) * poly(k)
      end do
    end do

    ! Right-biased WENO (u_{i+1/2}^+)
    call compute_weight_derivative(var, i, wgts, dwgt_plus_du, biased='right_stencil')
    poly(1) = dot_product(p1, var(i+3:i+1:-1))
    poly(2) = dot_product(p2, var(i+2:i  :-1))
    poly(3) = dot_product(p3, var(i+1:i-1:-1))
    du_plus_du = 0.0_r8
    do k = 1, 3
      do m = -2, 2
        j = i + m
        if (j < 1) j = j + nx 
        if (j > nx) j = j - nx
        du_plus_du(m) = du_plus_du(m) + wgts(k) * (&
           merge(p1(1), 0.0_r8, k==1 .and. m==-2) + & ! ∂p1/∂u_{i+3}
           merge(p1(2), 0.0_r8, k==1 .and. m==-1) + & ! ∂p1/∂u_{i+2}
           merge(p1(3), 0.0_r8, k==1 .and. m== 0) + & ! ∂p1/∂u_{i+1}
           merge(p2(1), 0.0_r8, k==2 .and. m==-1) + & ! ∂p2/∂u_{i+2}
           merge(p2(2), 0.0_r8, k==2 .and. m==0 ) + & ! ∂p2/∂u_{i+1}
           merge(p2(3), 0.0_r8, k==2 .and. m==1 ) + & ! ∂p2/∂u_{i}
           merge(p3(1), 0.0_r8, k==3 .and. m==0 ) + & ! ∂p3/∂u_{i+1}
           merge(p3(2), 0.0_r8, k==3 .and. m==1 ) + & ! ∂p3/∂u_{i}
           merge(p3(3), 0.0_r8, k==3 .and. m==2 )   & ! ∂p3/∂u_{i-1}
            ) + dwgt_plus_du(k,m) * poly(k)
      end do
    end do

    if (Riemann_solver == 'LF') then
      dF_du = 0.5_r8 * (u * (du_minus_du + du_plus_du) - dx / dt * (du_plus_du - du_minus_du))
    else if (Riemann_solver == 'LLF') then ! LLF flux Jacobian
      dF_du = 0.5_r8 * (u * (du_minus_du + du_plus_du) - alpha * (du_plus_du - du_minus_du))
    end if

  end subroutine compute_dF_du

  subroutine compute_weight_derivative(var, i, wgt, dwgt_dvar, biased)
    !===========================================================
    ! Compute WENO weights and their Jacobian ∂w_k/∂u_j
    ! wgt_k = alpha_k / sum(alpha_k), alpha_k = gamma_k / (eps + beta_k)^2
    ! chain rule: ∂wgt_k / ∂var = ∂wgt_k / ∂alpha_k * ∂alpha_k / ∂beta_k * ∂beta_k / ∂var
    !===========================================================
    real(r8), intent(in) :: var(1-nhalo:nx+nhalo)
    integer, intent(in) :: i
    character(len=*), intent(in) :: biased
    real(r8), intent(out) :: wgt(1:3), dwgt_dvar(1:3, -2:2)
    real(r8) :: beta(1:3), alpha(1:3), sum_alpha
    real(r8) :: dalpha_dbeta(1:3), dwgt_dalpha(1:3, 1:3), dbeta_du(1:3, -2:2)
    integer :: k, m

    ! Compute β_k and  𝛂_k
    call smoothness_indicators(var, i, beta, biased)
    alpha = gamma / (epsilon + beta)**2
    sum_alpha = sum(alpha)
    wgt = alpha / sum_alpha

    ! Compute ∂w_k/∂𝛂_m
    do k = 1, 3
      do m = 1, 3
        if (k == m) then
          dwgt_dalpha(k, m) = (sum_alpha - alpha(k)) / sum_alpha**2
        else
          dwgt_dalpha(k, m) = -alpha(k) / sum_alpha**2
        end if
      end do
    end do

    ! Compute ∂α_k/∂β_k
    dalpha_dbeta(1:3) = -2.0_r8 * gamma(1:3) / (epsilon + beta(1:3))**3

    if (trim(biased) == 'left_stencil') then
      ! Compute ∂β_k/∂u_j for all stencils
      dbeta_du = 0.0_r8
      ! Stencil S1 (i-2, i-1, i)
      dbeta_du(1, -2) = (13.0_r8 / 6.0_r8) * dot_product(s11, var(i-2:i)) + &
                        (1.0_r8  / 2.0_r8) * dot_product(s12, var(i-2:i))
      dbeta_du(1, -1) = (-13.0_r8 / 3.0_r8) * dot_product(s11, var(i-2:i)) - &
                                     2.0_r8 * dot_product(s12, var(i-2:i))
      dbeta_du(1,  0) = (13.0_r8 / 6.0_r8) * dot_product(s11, var(i-2:i)) + &
                        (3.0_r8  / 2.0_r8) * dot_product(s12, var(i-2:i))

      ! Stencil S2 (i-1, i, i+1)
      dbeta_du(2, -1) = (13.0_r8 / 6.0_r8) * dot_product(s21, var(i-1:i+1)) + &
                        ( 1.0_r8 / 2.0_r8) * dot_product(s22, var(i-1:i+1))
      dbeta_du(2,  0) = (-13.0_r8 / 3.0_r8) * dot_product(s21, var(i-1:i+1))
      dbeta_du(2,  1) = (13.0_r8 / 6.0_r8) * dot_product(s21, var(i-1:i+1)) + &
                        (-1.0_r8 / 2.0_r8) * dot_product(s22, var(i-1:i+1))

      ! Stencil S3 (i, i+1, i+2)
      dbeta_du(3,  0) = (13.0_r8 / 6.0_r8) * dot_product(s31, var(i:i+2)) + &
                        ( 3.0_r8 / 2.0_r8) * dot_product(s32, var(i:i+2))
      dbeta_du(3,  1) = (-13.0_r8/ 3.0_r8) * dot_product(s31, var(i:i+2)) - &
                                    2.0_r8 * dot_product(s32, var(i:i+2))
      dbeta_du(3,  2) = (13.0_r8 / 6.0_r8) * dot_product(s31, var(i:i+2)) + &
                         (1.0_r8 / 2.0_r8) * dot_product(s32, var(i:i+2))

    else if (trim(biased) == 'right_stencil') then
      ! Compute ∂β_k/∂u_j for all stencils
      dbeta_du = 0.0_r8
      ! Stencil S1 (i+3, i+2, i+1)
      dbeta_du(1, -2) = (13.0_r8 / 6.0_r8) * dot_product(s11, var(i+3:i+1:-1)) + &
                        (1.0_r8  / 2.0_r8) * dot_product(s12, var(i+3:i+1:-1))
      dbeta_du(1, -1) = (-13.0_r8 / 3.0_r8) * dot_product(s11, var(i+3:i+1:-1)) - &
                                     2.0_r8 * dot_product(s12, var(i+3:i+1:-1))
      dbeta_du(1,  0) = (13.0_r8 / 6.0_r8) * dot_product(s11, var(i+3:i+1:-1)) + &
                        (3.0_r8  / 2.0_r8) * dot_product(s12, var(i+3:i+1:-1))

      ! Stencil S2 (i+2, i+1, i)
      dbeta_du(2, -1) = (13.0_r8 / 6.0_r8) * dot_product(s21, var(i+2:i:-1)) + &
                        ( 1.0_r8 / 2.0_r8) * dot_product(s22, var(i+2:i:-1))
      dbeta_du(2,  0) = (-13.0_r8 / 3.0_r8) * dot_product(s21, var(i+2:i:-1))
      dbeta_du(2,  1) = (13.0_r8 / 6.0_r8) * dot_product(s21, var(i+2:i:-1)) + &
                        (-1.0_r8 / 2.0_r8) * dot_product(s22, var(i+2:i:-1))

      ! Stencil S3 (i+1, i, i-1)
      dbeta_du(3,  0) = (13.0_r8 / 6.0_r8) * dot_product(s31, var(i+1:i-1:-1)) + &
                        ( 3.0_r8 / 2.0_r8) * dot_product(s32, var(i+1:i-1:-1))
      dbeta_du(3,  1) = (-13.0_r8/ 3.0_r8) * dot_product(s31, var(i+1:i-1:-1)) - &
                                    2.0_r8 * dot_product(s32, var(i+1:i-1:-1))
      dbeta_du(3,  2) = (13.0_r8 / 6.0_r8) * dot_product(s31, var(i+1:i-1:-1)) + &
                         (1.0_r8 / 2.0_r8) * dot_product(s32, var(i+1:i-1:-1))
    end if

    ! Chain rule: ∂w_k/∂u_j = ∑_m (∂w_k/∂α_m * ∂α_m/∂β_m * ∂β_m/∂u_j)
    dwgt_dvar = 0.0_r8
    do k = 1, 3
      do m = 1, 3
        dwgt_dvar(k,-2:2) = dwgt_dvar(k,-2:2) + dwgt_dalpha(k,m) * dalpha_dbeta(m) * dbeta_du(m,-2:2)
      end do
    end do

  end subroutine compute_weight_derivative

  subroutine smoothness_indicators(var, i, beta, biased)
    ! Compute smoothness indicators β_k for stencil k
    real(r8), intent(in) :: var(1-nhalo:nx+nhalo)
    integer, intent(in) :: i
    character(len=*), intent(in) :: biased
    real(r8), intent(out) :: beta(1:3)
    
    if (trim(biased) == 'left_stencil') then
      beta(1) = s1 * dot_product(s11, var(i-2:i  ))**2 + s2 * dot_product(s12, var(i-2:i  ))**2
      beta(2) = s1 * dot_product(s21, var(i-1:i+1))**2 + s2 * dot_product(s22, var(i-1:i+1))**2
      beta(3) = s1 * dot_product(s31, var(i  :i+2))**2 + s2 * dot_product(s32, var(i  :i+2))**2
    else if (trim(biased) == 'right_stencil') then
      beta(1) = s1 * dot_product(s11, var(i+3:i+1:-1))**2 + s2 * dot_product(s12, var(i+3:i+1:-1))**2
      beta(2) = s1 * dot_product(s21, var(i+2:i  :-1))**2 + s2 * dot_product(s22, var(i+2:i  :-1))**2
      beta(3) = s1 * dot_product(s31, var(i+1:i-1:-1))**2 + s2 * dot_product(s32, var(i+1:i-1:-1))**2
    end if

  end subroutine smoothness_indicators

  subroutine compute_interface_flux_weno5(var, dt, flux)

    real(r8), intent(in) :: var(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    real(r8), intent(out) :: flux(0:nx)
    real(r8) :: weights(3), poly(3), beta(3), var_left, var_right
    real(r8) :: alpha
    integer i

    alpha = abs(u)
    do i = 1, nx
      poly(1) = dot_product(p1, var(i-2:i  ))
      poly(2) = dot_product(p2, var(i-1:i+1))
      poly(3) = dot_product(p3, var(i  :i+2))
      call smoothness_indicators(var, i, beta, biased='left_stencil')
      weights =  gamma / (epsilon + beta)**2
	    weights = weights / sum(weights)
	    var_left = dot_product(weights, poly)

      poly(1) = dot_product(p1, var(i+3:i+1:-1))
      poly(2) = dot_product(p2, var(i+2:i  :-1))
      poly(3) = dot_product(p3, var(i+1:i-1:-1))
      call smoothness_indicators(var, i, beta, biased='right_stencil')
      weights =  gamma / (epsilon + beta)**2
      weights = weights / sum(weights)
	    var_right = dot_product(weights, poly)

      if (Riemann_solver == 'LF') then ! Lax-Friedrichs flux
	      flux(i) = 0.5_r8 * (u * (var_right + var_left) - dx / dt * (var_right - var_left))
	    else if (Riemann_solver == 'LLF') then ! local Lax-Friedrichs flux
	      flux(i) = 0.5_r8 * (u * (var_right + var_left) - alpha * (var_right - var_left))
      end if 
    end do
    flux(0) = flux(nx)! periodic condition

  end subroutine compute_interface_flux_weno5

end module weno_mod